package com.yiwenup.struct._02_tree.bst;

import com.yiwenup.struct._02_tree.base.BinaryTree;

import java.util.Comparator;

/**
 * 二叉搜索树
 */
@SuppressWarnings("unchecked")
public class BinarySearchTree<E> extends BinaryTree<E> {

    private Comparator<E> comparator;

    public BinarySearchTree() {
        this(null);
    }

    public BinarySearchTree(Comparator<E> comparator) {
        this.comparator = comparator;
    }

    public void add(E element) {
        elementNotNullCheck(element);

        if (root == null) {
            root = createNode(element, null);
            size++;
            afterAdd(root);
            return;
        }

        // 找到父节点
        Node<E> parent = root;
        Node<E> node = root;
        int compare = 0;
        while (node != null) {
            compare = compare(element, node.element);
            parent = node;
            if (compare > 0) {
                node = node.right;
            } else if (compare < 0) {
                node = node.left;
            } else {
                node.element = element;
                return;
            }
        }

        //
        Node<E> newNode = createNode(element, parent);
        if (compare > 0) {
            parent.right = newNode;
        } else {
            parent.left = newNode;
        }
        size++;
        afterAdd(newNode);
    }

    protected void afterAdd(Node<E> node) {}

    protected void afterRemove(Node<E> node) {}

    public void remove(E element) {
        remove(node(element));
    }

    private void remove(Node<E> node) {
        if (node == null) return;

        if (node.hasTwoChildren()) {
            // 度为2
            Node<E> successor = successor(node);
            // 后继节点的值覆盖度为2的值
            node.element = successor.element;
            // 删除后继节点
            node = successor;
        }

        // 删除度为1或0的节点
        Node<E> replace = node.left != null ? node.left : node.right;

        if (replace != null) {
            // 度为1
            replace.parent = node.parent;
            if (node.parent == null) {
                // 度为1，且是根节点
                root = replace;
            } else if (node == node.parent.left) {
                node.parent.left = replace;
            } else {
                node.parent.right = replace;
            }

            afterRemove(replace);
        } else if (node.parent == null) {
            // 度为0，且是根节点
            root = null;
            afterRemove(node);
        } else {
            // 度为0，但不是根节点 -> 叶子节点
            if (node.parent.left == node) {
                node.parent.left = null;
            } else {
                node.parent.right = null;
            }
            afterRemove(node);
        }

        size--;
    }

    public boolean contains(E element) {
        return node(element) != null;
    }

    private void elementNotNullCheck(E element) {
        if (element == null) {
            throw new IllegalArgumentException("element must not be null");
        }
    }

    /**
     * @return 0 -> e1==e2; <0 -> e1<e2; >0 -> e1>e2
     */
    private int compare(E e1, E e2) {
        if (comparator != null) {
            return comparator.compare(e1, e2);
        }
        return ((Comparable<E>) e1).compareTo(e2);
    }


    /**
     * 根据元素查询节点
     *
     * @param element 元素
     * @return 节点
     */
    private Node<E> node(E element) {
        Node<E> node = this.root;
        while (node != null) {
            int compare = compare(element, node.element);
            if (compare == 0) return node;
            if (compare > 0) {
                node = node.right;
            } else {
                node = node.left;
            }
        }
        return null;
    }


}
